import torch
import os
import cv2
from sklearn.cluster import dbscan
from sklearn import mixture
import matplotlib.pyplot as plt
import numpy as np


def div_lines(
    img: torch.TensorType,
    is_show: bool = False,
    cache_dir: os.PathLike = ".",
    is_force: bool = True,
    img_name: str = "",
) -> tuple[list[dict], plt.Axes]:
    if not os.path.exists(os.path.join(cache_dir, f"lines_{img_name}.pt")) or is_force:
        img = 255 - img
        bool_mask = torch.zeros_like(img, device=img.device)
        bool_mask[img > 30] = 255
        X = torch.argwhere(bool_mask > 0).cpu().numpy()
        _, idx = dbscan(X, eps=4, min_samples=2)
        num_ind = idx.max()
        lines = []  # 存放输出线条点集
        for i in range(num_ind):
            X_box = X[idx == i, :]
            if X_box.shape[0] > 20000 or X_box.shape[0] < 300:
                continue
            try:
                GMModel = mixture.GaussianMixture(
                    n_components=6, tol=1e-6, max_iter=500
                ).fit(X_box)
                idx2 = GMModel.predict(X_box)
                temp = {}
                temp["box"] = X_box
                temp["idx"] = idx2
                temp["valid"] = True
                lines.append(temp)
            except Exception as exc:
                print(exc)
                temp = {}
                temp["valid"] = False
                lines.append(temp)
                continue
        torch.save(lines, f"lines_{img_name}.pt")
    else:
        lines = torch.load(f"lines_{img_name}.pt")
    if is_show:
        axe = plt.subplot(111)
        j = 0
        for line in lines:
            if line["valid"]:
                for i in range(6):
                    line1 = line["box"][line["idx"] == i, :]
                    axe.scatter(
                        line1[:, 1],
                        900 - line1[:, 0],
                        s=np.ones_like(line1[:, 1]) * i,
                    )
                    axe.text(
                        np.mean(line1[:, 1]),
                        900 - np.mean(line1[:, 0]),
                        str(i),
                        bbox=dict(facecolor="orange", alpha=0.5),
                    )
                axe.text(
                    np.mean(line["box"][:, 1]) + 20,
                    np.mean(900 - line["box"][:, 0]) + 20,
                    str(j),
                    bbox=dict(facecolor="blue", alpha=0.2),
                )
                j += 1
    return lines, axe


def esti_lines(lines: list[dict], line_id: list[list], **kwargs) -> tuple[list, list]:
    L1 = []
    L2 = []
    for id in line_id:
        line1 = lines[id[0]]["box"][lines[id[0]]["idx"] == id[1], :]
        line2 = lines[id[0]]["box"][lines[id[0]]["idx"] == id[2], :]
        L1.append(
            main_esti(
                np.pad(line1, ((0, 0), (0, 1)), mode="constant", constant_values=1),
                **kwargs,
            ).tolist()
        )
        L2.append(
            main_esti(
                np.pad(line2, ((0, 0), (0, 1)), mode="constant", constant_values=1),
                **kwargs,
            ).tolist()
        )
    return L1, L2


def main_esti(
    line: np.ndarray, axe: plt.Axes = None, img: np.ndarray = None, is_show: bool = True
):
    _, _, Vh = np.linalg.svd(line)
    l = Vh[-1, :].T
    x1 = np.arange(line[:, 0].min(), line[:, 0].max(), 0.1)
    y1 = -(l[0] * x1 + l[2]) / l[1]
    len_1 = y1.max() - y1.min()
    y2 = np.arange(line[:, 1].min(), line[:, 1].max(), 0.1)
    x2 = -(l[1] * y2 + l[2]) / l[0]
    len_2 = x2.max() - x2.min()
    if len_1 > len_2:
        x = x2
        y = y2
    else:
        x = x1
        y = y1
    if axe == None:
        plt.plot(y, 900 - x, "bo")
        if is_show:
            img = cv2.line(
                img,
                (int(y[0]), int(x[0])),
                (int(y[-1]), int(x[-1])),
                (0, 0, 255),
                thickness=2,
            )
    else:
        if is_show:
            cv2.line(
                img,
                (int(y[0]), int(x[0])),
                (int(y[-1]), int(x[-1])),
                (0, 0, 255),
                thickness=2,
            )
        axe.plot(y, 900 - x, "bo")
    return l


def preprocess(img: np.ndarray, is_show: bool = False):
    img1 = np.mean(img, axis=2)
    img1 = (img1 - np.min(img1)) / (np.max(img1) - np.min(img1)) * 255
    img1 = img1.astype(np.uint8)
    cv2.medianBlur(img1, 5, img1)
    cv2.GaussianBlur(img1, (5, 5), 1, img1)
    return torch.tensor(img1, device="cuda")


def extract_main(
    line_id: list,
    *,
    output_dir: str = "./output",
    is_force: bool = True,
    is_show: bool = True,
    img_name: str = "img_project",
) -> tuple[list[list], np.ndarray]:
    # load the background pictures
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    # load the pictures to be rectified
    # img = cv2.imread(os.path.join(output_dir, f"img_project{postfix}.png"))
    img = cv2.imread(os.path.join(output_dir, img_name + ".png"))
    # img_1 = torch.tensor(np.mean(img, axis=2), device="cuda")
    img_1 = preprocess(img, is_show=is_show)
    # extract the lines
    lines, axe = div_lines(
        img_1,
        is_show=True,
        is_force=is_force,
        img_name=img_name,
    )
    # Using five orgonally crossing lines to rectify projective transformation
    # Choose five lines, using box index and line index respectively, each row represent one crossing lines, using format like [box_id,line1_id,line2_id;  ...]
    # estimate the lines
    L1, L2 = esti_lines(lines, line_id, axe=axe, img=img, is_show=is_show)
    plt.savefig(os.path.join(output_dir, f"lines_{img_name}.png"))
    plt.cla()

    return [[*l1, *l2] for l1, l2 in zip(L1, L2)], torch.tensor(img, device="cuda")
